1754E - Wish I Knew How to Sort - CodeForces Solution

#include "bits/stdc++.h"
#define int long long
using namespace std;
const int N=2e5+2,mod=998244353;
int T,n,s,k,b,a[N],inv[N],ans;
void solve(){
    scanf("%lld",&n);b=n*(n-1ll)/2%mod,s=k=ans=0;
    for (int i = 1; i <=n ; ++i) scanf("%lld",&a[i]),s+=a[i];
    for (int i = 1; i <=n-s ; ++i) k+=a[i];
    for (int i = k; i ; i--) ans=(ans+b*inv[i]%mod*inv[i])%mod;
    printf("%lld\n",ans);
}
signed main(){
    inv[1]=1;
    for(int i=2;i<N;++i) inv[i]=1ll*(mod-mod/i)*inv[mod%i]%mod;
    scanf("%lld",&T);
    while (T--){
        solve();
    }
    return 0;
}